# 各类建筑的父类
from const import *


class Construct(object):
    def __init__(self, serial=0, name='建筑物', constructtype='单向闸', upserial=0, downserial=1, baselevel=0, width=1,
                 hight=1, maxflow=100, rule=None):
        # rule = [上游水位上限，上游水位下限,下游水位上限，下游水位下限]    上下游水位在此范围内方可开闸/机器
        # 默认是常开，即rule=[[float('inf'), float('-inf')], [float('inf'), float('-inf')]]

        if rule is None:
            rule = [float('inf'), float('-inf'), float('inf'), float('-inf')]
        self.serial = serial  # 每个建筑均应有序号
        self.name = name  # 建筑物名称
        self.constructtype = constructtype
        # 需对上下游序号进行检查，以避免出现计算顺序错误
        if upserial >= downserial and constructtype in ['双向闸', '单向闸']:
            raise ValueError('上游序号大于等于下游序号，将导致计算错误')
        self.upserial = upserial  # 每个建筑物上下游均应有储水功能
        self.downserial = downserial
        self.baselevel = baselevel  # 底板高程
        self.width = width
        self.hight = hight
        self.maxflow = maxflow  # 请注意这个是单向的还是双向的
        self.statelist = []  # 用于存储开关状态
        self.flowlist = []  # 用于计算过闸流量
        self.rule = rule  # 仅当在给定的范围中，该建筑物才会开启
        if self.constructtype in ['单向闸', '泵站']:
            self.flowmethod = 'single'
        elif self.constructtype == '双向闸':
            self.flowmethod = 'both'
        # self.folw = 0    # 设置默认流量   应当置于后文
        # self.state = False     # 设置状态
        self._checkrule()
        pass

    def flow_calculation(self, upwaterlevel, downwaterlevel, steptime):
        self._checkrule()
        # print(self.rule, upwaterlevel, downwaterlevel)
        # 根据上下游水位判断是否开启闸门，并计算流量
        if self.rule[0] > upwaterlevel > self.rule[1] and self.rule[2] > downwaterlevel > self.rule[3]:
            # 当满足条件时开启闸门，并采用堰流公式计算流量
            self.statelist.append(True)
            flow = (upwaterlevel - downwaterlevel) * self.width * self.hight * steptime
            if flow > self.maxflow:
                # 当大于最大入流时设置为最大入流
                flow = self.maxflow
            else:
                if flow < 0 and self.flowmethod == 'single':
                    # 当流量小于0时，且入流为单向时流量设置为0、当前闸门设置为关闭
                    flow = 0
                    self.statelist[-1] = False
                elif flow < -self.maxflow:
                    # 当流量小于0，且入流设置为双向，则受最大流量限制
                    flow = -self.maxflow
            self.flowlist.append(flow)
            # print('%s号闸门当前时段开启' % self.serial)
            # print('当前计算流量为%s' % flow)
            return flow
        else:
            self.statelist.append(False)
            self.flowlist.append(0)
            return 0
    def reflow(self, reward_flow):
        # 减少当前建筑物的出流，以满足水量平衡
        if reward_flow > 0:
            raise ValueError('水量平衡有误，请检查代码')
        elif reward_flow == 0:
            return reward_flow
        else:
            flow = reward_flow + self.flowlist[-1]
            if self.flowmethod == 'single':
                flow = max(0, flow)    # 当为单向流时，确保流量不低于0
            elif self.flowmethod == 'both':
                flow = max(-self.maxflow, flow)  # 当为双向流时，确保流量不高于最大流量
            reward_flow = self.flowlist[-1] - flow + reward_flow
            # 尚缺值 = （原计划流量 - 更正后流量）【即变更后新增流量】 + 原缺值
            self.flowlist[-1] = flow
            return reward_flow

    def _checkrule(self):
        if self.rule == []:
            self.rule = rule0


class pump(Construct):
    def __init__(self, serial=0, name='泵站1', constructtype='泵站', upserial=0, downserial=1, width=1, hight=1,
                 maxflow=100, rule=[float('inf'), float('-inf'), float('inf'), float('-inf')],
                 target: list = None):
        super().__init__(serial=serial, name=name, constructtype=constructtype, upserial=upserial,
                         downserial=downserial,
                         width=width, hight=hight, maxflow=maxflow, rule=rule)
        self.baselevel = []
        if target is None:
            self.target = [self.maxflow]
        else:
            self.target = target

    def flow_calculation(self, upwaterlevel, downwaterlevel, steptime):
        # 根据上下游水位判断是否开启闸门，并计算流量
        if self.rule[0] > upwaterlevel > self.rule[1] and self.rule[2] > downwaterlevel > self.rule[3]:
            # 当满足条件时开启闸门，并以目标流量作为其流量
            # 根据当前流量系列长度，结合目标流量系列，确定目标流量
            flowtarget = self.target[len(self.flowlist) % len(self.target)]
            self.statelist.append(True)
            flow = flowtarget * steptime
            self.flowlist.append(flow)
        else:
            self.statelist.append(False)
            flow = 0
            self.flowlist.append(flow)
        return flow
